from omegaconf import DictConfig

from lightning import Trainer
from lightning.pytorch.loggers import WandbLogger

from ecgcmr.utils.misc import fix_seed
from ecgcmr.signal.sig_datasets.MaskedECGLightning import MaskedECGDataModule
from ecgcmr.signal.sig_datasets.DownstreamECGLightning import DownstreamECGDataModule
from ecgcmr.signal.sig_models.ECGViTMAE import ECGViTMAE
from ecgcmr.signal.sig_models.ECGViTEval import ECGViTEval


def train_ecg_masked(
        cfg: DictConfig,
        wandb_logger: WandbLogger,
        save_dir: str,
        devices: int = 1
        ):
    fix_seed(seed=cfg.seed)
    
    datamodule = MaskedECGDataModule(cfg=cfg)
    
    model = ECGViTMAE(cfg=cfg, save_dir=save_dir)
    wandb_logger.watch(model, log_graph=False)

    strategy = "ddp" if devices > 1 else "auto"

    trainer = Trainer(
        accelerator="gpu",
        devices=devices,
        strategy=strategy,
        precision="bf16-mixed",
        logger=wandb_logger,
        max_epochs=cfg.max_epochs,
        log_every_n_steps=cfg.log_every_n_steps,
        check_val_every_n_epoch=cfg.check_val_every_n_epoch,
        default_root_dir=save_dir,
        num_sanity_val_steps=0,
        profiler="simple"
        )
    
    trainer.fit(model=model, datamodule=datamodule)

def finetune_ecg_masked(cfg: DictConfig,
                        wandb_logger: WandbLogger,
                        save_dir: str = None,
                        devices: int = 1,
                        checkpoint_path = None):
    
    datamodule = DownstreamECGDataModule(cfg=cfg)

    test_model = ECGViTEval(cfg=cfg, save_dir=save_dir, pretrained_model_name_or_path=checkpoint_path)

    wandb_logger.watch(test_model, log_graph=False)

    trainer = Trainer(
        accelerator="gpu",
        devices=devices,
        precision="bf16-mixed",
        logger=wandb_logger,
        max_epochs=cfg.downstream_task.max_epochs,
        log_every_n_steps=cfg.log_every_n_steps,
        default_root_dir=save_dir,
        num_sanity_val_steps=0,
    )

    trainer.fit(model=test_model, datamodule=datamodule)